import os
import time
import json
import pickle
import logging

import torch
import cv2
import numpy as np
from easydict import EasyDict

from vima_bench import *
from gym.wrappers import TimeLimit as _TimeLimit
from gym import Wrapper

from request_post import PostAPI


def exception_handler(exception, logger, **kwargs):
    logger.error("Exception: {}".format(exception))
    task_info = {
        "task_id": kwargs["task_id"],
        "task": kwargs["whole_task"],
        "exec": kwargs["exec_codes"],
        "skip": True,
        "success": False,
        "exception": str(exception),
    }
    return task_info


class ResetFaultToleranceWrapper(Wrapper):
    max_retries = 10

    def __init__(self, env):
        super().__init__(env)

    def reset(self):
        for _ in range(self.max_retries):
            try:
                return self.env.reset()
            except:
                current_seed = self.env.unwrapped.task.seed
                self.env.global_seed = current_seed + 1
        raise RuntimeError(
            "Failed to reset environment after {} retries".format(self.max_retries)
        )


class TimeLimitWrapper(_TimeLimit):
    def __init__(self, env, bonus_steps: int = 0):
        super().__init__(env, env.task.oracle_max_steps + bonus_steps)


@torch.no_grad()
def main(cfg):
    import logging
    logging.basicConfig(level=logging.DEBUG,  # 设置日志级别
                    format="%(asctime)s - %(levelname)s - %(message)s") 
    logging.info("cfg: {}".format(cfg))
    debug_flag = cfg.debug_flag
    # assert cfg.partition in ALL_PARTITIONS
    # assert cfg.task in PARTITION_TO_SPECS["test"][cfg.partition]

    seed = cfg.seed
    env = TimeLimitWrapper(
        ResetFaultToleranceWrapper(
            make(
                cfg.task,
                modalities=["segm", "rgb"],
                task_kwargs=PARTITION_TO_SPECS["test"][cfg.partition][cfg.task],
                seed=seed,
                render_prompt=False,
                display_debug_window=debug_flag,
                hide_arm_rgb=False,
            )
        ),
        bonus_steps=2,
    )
    single_model_flag = True if cfg.prompt_modal == "single" else False
    result_folder = (
        cfg.save_dir + "/" + cfg.partition + "/" + cfg.task + "/" + cfg.prompt_modal
    )
    if not os.path.exists(result_folder):
        os.makedirs(result_folder)
    eval_res_name = cfg.partition + "_" + cfg.task + ".json"
    eval_result_file_path = os.path.join(result_folder, eval_res_name)

    task_id = 0
    all_infos = []
    if cfg.reuse and os.path.exists(eval_result_file_path):
        with open(eval_result_file_path, "r") as f:
           all_infos = json.load(f)
    
    while True:
        try:
            env.global_seed = seed

            obs = env.reset()
            env.render()
            meta_info = env.meta_info
            prompt = env.prompt
            prompt_assets = env.prompt_assets
            if cfg.task == "visual_manipulation":
                dragged_obj = prompt_assets["dragged_obj_1"]["segm"]["obj_info"]["obj_color"] + " " + prompt_assets["dragged_obj_1"]["segm"]["obj_info"]["obj_name"]
                base_obj = prompt_assets["base_obj"]["segm"]["obj_info"]["obj_color"] + " " + prompt_assets["base_obj"]["segm"]["obj_info"]["obj_name"]
                instruction = prompt.format(dragged_obj_1=dragged_obj, base_obj=base_obj)
            elif cfg.task == "manipulate_old_neighbor":
                dragged_obj = prompt_assets["dragged_obj"]["segm"]["obj_info"]["obj_color"] + " " + prompt_assets["dragged_obj"]["segm"]["obj_info"]["obj_name"]
                base_obj = prompt_assets["base_obj"]["segm"]["obj_info"]["obj_color"] + " " + prompt_assets["base_obj"]["segm"]["obj_info"]["obj_name"]
                instruction = prompt.format(dragged_obj=dragged_obj, base_obj=base_obj)
            else:
                instruction = prompt.replace(" in {scene}", "")
            
            print(instruction)
            # base_obj_2 = prompt_assets["base_obj_2"]["segm"]["obj_info"]["obj_color"] + " " + prompt_assets["base_obj_2"]["segm"]["obj_info"]["obj_name"]
            # instruction = instruction.format(base_obj=base_obj)


            # print(instruction)
            # instruction = "Please rearrange the objects to the target state."
            task_id += 1
            logging.info(f"==================Task {task_id}=========================")

            if task_id >= 150:
                break
            done = False
            elapsed_steps = 0
            IMAGE_INIT_top = np.transpose(obs["rgb"]["top"], (1, 2, 0))

            while True:
                from PIL import Image
                top_img = np.transpose(obs["rgb"]["top"], (1, 2, 0))
                image = Image.fromarray(top_img)
                time.sleep(10)
                image.save("top_img.png")
                actions = post_api.post_request(init_image=top_img, prompt=instruction)
                
                for action in actions:
                    print(action)
                    obs, _, done, info = env.step(action)
                elapsed_steps += 1

                if done and info:
                    task_info = {
                        "task_id": task_id,
                        "steps": elapsed_steps,
                        "success": info["success"],
                        "failure": info["failure"],
                    }
                else:
                    task_info = {
                        "task_id": task_id,
                        "steps": elapsed_steps,
                        "success": False,
                        "failure": False,
                    }

                logging.info(
                    f"task id: {task_info['task_id']} success: {task_info['success']}"
                )
                
                if cfg.reuse and task_id - 1 < len(all_infos):
                    all_infos[task_id - 1] = task_info

                all_infos.append(task_info)
                with open(eval_result_file_path, "w") as f:
                    json.dump(all_infos, f)
              
                break

        except Exception as e:
            logging.error(f"Error occurred: {e}")
            continue

    success_rate = sum([info["success"] for info in all_infos]) / len(all_infos)
    logging.warning(msg="==================Evaluation Done=========================")
    logging.info(cfg)
    logging.info("Success rate: {}".format(success_rate))
    env.env.close()
    del env

if __name__ == "__main__":
    post_api = PostAPI(
        url = "http://127.0.0.1:8080/qwen2_vl_service/chat",
        remote_ip="connect.bjc1.seetacloud.com",
        remote_user_name="root",
        port=49371,
        remote_password="8qm9vWezztWs",
        remote_path="/root/autodl-fs/qwen2_vl_service/images",
        local_path="/Users/yangzhikai/Documents/code/robovlm/local_images"
    ) 
    prompt_modal = ["withput_cot"]
    # prompt_modal = ["multi", "single"]
    tasks = [
        # "visual_manipulation",
        # "rotate",
        # "pick_in_order_then_restore",
        # "rearrange_then_restore",
        # "rearrange",
        "manipulate_old_neighbor",
        # "scene_understanding"
    ]
    partitions = [
        "placement_generalization",
        # "combinatorial_generalization",
        # "novel_object_generalization"
    ]
    save_dir = "save_demo_output"
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    seed = 42
    hide_arm =False # False for demo usage, True for eval usage
    for task in tasks:
        for partition in partitions:
            for modal in prompt_modal:
                eval_cfg = {
                    "partition": partition,
                    "task": task,
                    "device": "cuda",
                    "prompt_modal": modal,
                    "reuse": False,
                    "save_dir": save_dir,
                    "debug_flag": True,
                    "hide_arm": hide_arm,
                    "seed": seed,
                }
                logger_file = (
                    save_dir
                    + "/eval_on_seed_{}_hide_arm_{}_{}_{}_{}_modal.log".format(
                        eval_cfg["seed"],
                        eval_cfg["hide_arm"],
                        partition,
                        task,
                        modal,
                    )
                )
                if os.path.exists(path=logger_file):
                    os.remove(logger_file)
                main(EasyDict(eval_cfg))
